import argparse
import os
import os.path as osp
import torch
from models.convenience import PCCSDataset, TrainMode
from utils.utility import (create_saved_folders, 
                           getLogger, 
                           acquire_train_mode)
from utils.checkers import check_dataset

# TODO: Define train steps
from train_steps import train


logger = getLogger(__name__)


# TODO: refine to fit this task
parser = argparse.ArgumentParser()
# Training associated
parser.add_argument("-train", "--train", action="store_true")
parser.add_argument("-eval", "--eval", action="store_true")
parser.add_argument("-test", "--test", action="store_true")
parser.add_argument("-tf", "--test_filename", type=str, nargs='+')
parser.add_argument("-o", "--operations", type=str, nargs="+")

# Dataset associated
parser.add_argument("-dsbp", "--dataset_base_path", type=str, default="OGM-datasets")
parser.add_argument("-dsn", "--dataset_name", type=str, default="OGM-Turtlebot2")

# Saving directories
parser.add_argument("-sd", "--save_dir", type=str, default="models")
parser.add_argument("-mp", "--model_save_path", type=str, default="model")
parser.add_argument("-sflag", "--save_flag", type=str, default="debug")

# TODO: data augmentation

# Model related settings
parser.add_argument("-c", "--n_cluster", type=int, default=200)
parser.add_argument("-bs", "--batch_size", type=int, default=128)
parser.add_argument("--epochs", type=int, default=100)
parser.add_argument("--base_lr", type=float, default=5e-4)
parser.add_argument("--obs_len", type=int, default=10)
parser.add_argument("--pred_len", type=int, default=10)
parser.add_argument("--obs_feat_channels", type=int, default=128)
parser.add_argument("--pred_feat_channels", type=int, default=128)
parser.add_argument("--encoder_layer", type=int, default=1)

# Training related settings
parser.add_argument("--print_every", type=int, default=1)
parser.add_argument("--rad", type=float, default=1.0, help="Radius for Modality Loss")
parser.add_argument("--retrain", action="store_true")

args = parser.parse_args()

def main(args):

    # Specify the operations
    if args.train:
        op_codes = list(range(5))
        args.operations = acquire_train_mode(TrainMode, op_codes)
        logger.info(f"Operations including {[item.name for item in args.operations]}")
    elif args.eval:
        op_codes = [5]
        args.operations = acquire_train_mode(TrainMode, op_codes)
        logger.info(f"Operations including {[item.name for item in args.operations]}")
    elif args.test:
        op_codes = [6]
        args.operations = acquire_train_mode(TrainMode, op_codes)
        logger.info(f"Operations include {[item.name for item in args.operations]}")

    if args.operations is None:
        logger.error("Please specify the training operations, ABORT")
        exit()
    
    if not args.test:
        dset_path = osp.join(args.dataset_base_path, args.dataset_name)
        if check_dataset(dset_path):
            logger.info(f"Train on dataset: {args.dataset_name}")
            logger.info(f"Dataset path: {dset_path}")
        train_set = PCCSDataset(dset_path, 'train')
        val_set = PCCSDataset(dset_path, 'val')
    else:
        train_set = None
        val_set = PCCSDataset(dset_path, 'test')

    # Create the directory for saving the model details
    save_model_dir = create_saved_folders(args.save_dir, args.dataset_name, args.save_flag)

    # The default device
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    for op in args.operations:
        train(train_set, val_set, op, args.dataset_name, args, save_model_dir, device)




if __name__ == "__main__":
    main(args)